# TODO: Add comment
# 
# Author: yaping
###############################################################################

library("gplots")

heatmap.3 <- function (x, Rowv = TRUE, Colv = if (symm) "Rowv" else TRUE, 
		distfun = dist, hclustfun = hclust, dendrogram = c("both", 
				"row", "column", "none"), symm = FALSE, scale = c("none", 
				"row", "column"), na.rm = TRUE, revC = identical(Colv, 
				"Rowv"), add.expr, breaks, symbreaks = min(x < 0, na.rm = TRUE) || 
				scale != "none", col = "heat.colors", colsep, rowsep, 
		sepcolor = "white", sepwidth = c(0.05, 0.05), cellnote, notecex = 1, 
		notecol = "cyan", na.color = par("bg"), trace = c("column", 
				"row", "both", "none"), tracecol = "cyan", hline = median(breaks), 
		vline = median(breaks), linecol = tracecol, margins = c(5, 
				5), ColSideColors=NULL, RowSideColors=NULL, cexRow = 0.2 + 1/log10(nr), RowSideColorsName=NULL,
		cexCol = 0.2 + 1/log10(nc), labRow = NULL, labCol = NULL, 
		key = TRUE, keysize = 1.5, keyName="Methylation", density.info = c("histogram", 
				"density", "none"), denscol = tracecol, symkey = min(x < 
						0, na.rm = TRUE) || symbreaks, densadj = 0.25, main = NULL, 
		xlab = NULL, ylab = NULL, lmat = NULL, lhei = NULL, lwid = NULL, 
		addAverage = FALSE, averageData=NULL,clusterNums=NULL, clusterOrder =NULL, averageLegend=FALSE, colAveragePlot=NULL,logScale=FALSE,
		xAxisSeqForAvePlot=NULL, yAxisSeqForAvePlot=NULL,y_max=100, y_min=0, axisSeq=NULL, dataStep=NULL,dataSeq=NULL,xlabForAvePlot="Distance to elements (bp)",ylabForAvePlot="Methylation",
		bin_size_in_align=1, multipleAveragePlot=FALSE, sampleLen=NULL,
		...) 
{
	scale01 <- function(x, low = min(x), high = max(x)) {
		x <- (x - low)/(high - low)
		x
	}
	retval <- list()
	scale <- if (symm && missing(scale)) 
				"none"
			else match.arg(scale)
	dendrogram <- match.arg(dendrogram)
	trace <- match.arg(trace)
	density.info <- match.arg(density.info)
	if (length(col) == 1 && is.character(col)) 
		col <- get(col, mode = "function")
	if (!missing(breaks) && (scale != "none")) 
		warning("Using scale=\"row\" or scale=\"column\" when breaks are", 
				"specified can produce unpredictable results.", "Please consider using only one or the other.")
	if (is.null(Rowv) || is.na(Rowv)) 
		Rowv <- FALSE
	if (is.null(Colv) || is.na(Colv)) 
		Colv <- FALSE
	else if (Colv == "Rowv" && !isTRUE(Rowv)) 
		Colv <- FALSE
	if (length(di <- dim(x)) != 2 || !is.numeric(x)) 
		stop("`x' must be a numeric matrix")
	nr <- di[1]
	nc <- di[2]
	if (nr <= 1 || nc <= 1) 
		stop("`x' must have at least 2 rows and 2 columns")
	if (!is.numeric(margins) || length(margins) != 2) 
		stop("`margins' must be a numeric vector of length 2")
	if (missing(cellnote)) 
		cellnote <- matrix("", ncol = ncol(x), nrow = nrow(x))
	if (!inherits(Rowv, "dendrogram")) {
		if (((!isTRUE(Rowv)) || (is.null(Rowv))) && (dendrogram %in% 
					c("both", "row"))) {
			if (is.logical(Colv) && (Colv)) 
				dendrogram <- "column"
			else dedrogram <- "none"
			warning("Discrepancy: Rowv is FALSE, while dendrogram is `", 
					dendrogram, "'. Omitting row dendogram.")
		}
	}
	if (!inherits(Colv, "dendrogram")) {
		if (((!isTRUE(Colv)) || (is.null(Colv))) && (dendrogram %in% 
					c("both", "column"))) {
			if (is.logical(Rowv) && (Rowv)) 
				dendrogram <- "row"
			else dendrogram <- "none"
			warning("Discrepancy: Colv is FALSE, while dendrogram is `", 
					dendrogram, "'. Omitting column dendogram.")
		}
	}
	if (inherits(Rowv, "dendrogram")) {
		ddr <- Rowv
		#rowInd <- order.dendrogram(ddr)
		rowInd <- 1:nr
	}
	else if (is.integer(Rowv)) {
		hcr <- hclustfun(distfun(x))
		ddr <- as.dendrogram(hcr)
		ddr <- reorder(ddr, Rowv)
		rowInd <- order.dendrogram(ddr)
		if (nr != length(rowInd)) 
			stop("row dendrogram ordering gave index of wrong length")
	}
	else if (isTRUE(Rowv)) {
		Rowv <- rowMeans(x, na.rm = na.rm)
		hcr <- hclustfun(distfun(x))
		ddr <- as.dendrogram(hcr)
		ddr <- reorder(ddr, Rowv)
		rowInd <- order.dendrogram(ddr)
		if (nr != length(rowInd)) 
			stop("row dendrogram ordering gave index of wrong length")
	}
	else {
		#rowInd <- nr:1
		rowInd <- 1:nr
	}
	if (inherits(Colv, "dendrogram")) {
		ddc <- Colv
		#colInd <- order.dendrogram(ddc)
		colInd <- 1:nc
	}
	else if (identical(Colv, "Rowv")) {
		if (nr != nc) 
			stop("Colv = \"Rowv\" but nrow(x) != ncol(x)")
		if (exists("ddr")) {
			ddc <- ddr
			colInd <- order.dendrogram(ddc)
		}
		else colInd <- rowInd
	}
	else if (is.integer(Colv)) {
		hcc <- hclustfun(distfun(if (symm) 
									x
								else t(x)))
		ddc <- as.dendrogram(hcc)
		ddc <- reorder(ddc, Colv)
		colInd <- order.dendrogram(ddc)
		if (nc != length(colInd)) 
			stop("column dendrogram ordering gave index of wrong length")
	}
	else if (isTRUE(Colv)) {
		Colv <- colMeans(x, na.rm = na.rm)
		hcc <- hclustfun(distfun(if (symm) 
									x
								else t(x)))
		ddc <- as.dendrogram(hcc)
		ddc <- reorder(ddc, Colv)
		colInd <- order.dendrogram(ddc)
		if (nc != length(colInd)) 
			stop("column dendrogram ordering gave index of wrong length")
	}
	else {
		colInd <- 1:nc
	}
	retval$rowInd <- rowInd
	retval$colInd <- colInd
	retval$call <- match.call()
	x <- x[rowInd, colInd]
	x.unscaled <- x
	cellnote <- cellnote[rowInd, colInd]
	if (is.null(labRow)) 
		labRow <- if (is.null(rownames(x))) 
					(1:nr)[rowInd]
				else rownames(x)
	else labRow <- labRow[rowInd]
	if (is.null(labCol)) 
		labCol <- if (is.null(colnames(x))) 
					(1:nc)[colInd]
				else colnames(x)
	else labCol <- labCol[colInd]
	if (scale == "row") {
		retval$rowMeans <- rm <- rowMeans(x, na.rm = na.rm)
		x <- sweep(x, 1, rm)
		retval$rowSDs <- sx <- apply(x, 1, sd, na.rm = na.rm)
		x <- sweep(x, 1, sx, "/")
	}
	else if (scale == "column") {
		retval$colMeans <- rm <- colMeans(x, na.rm = na.rm)
		x <- sweep(x, 2, rm)
		retval$colSDs <- sx <- apply(x, 2, sd, na.rm = na.rm)
		x <- sweep(x, 2, sx, "/")
	}
	if (missing(breaks) || is.null(breaks) || length(breaks) < 
			1) {
		if (missing(col) || is.function(col)) 
			breaks <- 16
		else breaks <- length(col) + 1
	}
	if (length(breaks) == 1) {
		if (!symbreaks) 
			breaks <- seq(min(x, na.rm = na.rm), max(x, na.rm = na.rm), 
					length = breaks)
		else {
			extreme <- max(abs(x), na.rm = TRUE)
			breaks <- seq(-extreme, extreme, length = breaks)
		}
	}
	nbr <- length(breaks)
	ncol <- length(breaks) - 1
	if (class(col) == "function") 
		col <- col(ncol)
	min.breaks <- min(breaks)
	max.breaks <- max(breaks)
	x[x < min.breaks] <- min.breaks
	x[x > max.breaks] <- max.breaks
	if (missing(lhei) || is.null(lhei)){
		lhei <- c(keysize, 4)
		if(addAverage){
			lhei <- c(keysize, 4, 2)
		}

	}
		#lhei <- c(keysize, 4)
	if (missing(lwid) || is.null(lwid)) 
		lwid <- c(keysize, 4)
	if (missing(lmat) || is.null(lmat)) {
		lmat <- rbind(4:3, 2:1)
		if(addAverage){
			lmat <- rbind(4:3, 2:1,c(6,5))
		}
		
		
		if (!is.null(ColSideColors)) { ##need to come back to fix colsideColor' display problem!!
			if (!is.matrix(ColSideColors)) 
				stop("'ColSideColors' must be a matrix")
			if (!is.character(ColSideColors) || dim(ColSideColors)[1] != 
					nc) 
				stop("'ColSideColors' dim()[2] must be a character vector of length ncol(x)")
			lmat <- rbind(lmat[1, ] + 1, c(NA, 1), lmat[2, ] + 
							1)
			lhei <- c(lhei[1], 0.2*length(colnames(ColSideColors)), lhei[2])
		}
		if (!is.null(RowSideColors)) { ##temporary usage
			if (!is.matrix(RowSideColors)) 
				stop("'RowSideColors' must be a matrix")
			if (!is.character(RowSideColors) || dim(RowSideColors)[1] != 
					nr) 
				stop("'RowSideColors' dim()[2] must be a character vector of length nrow(x)")
			
			
			if(addAverage){
				lmat <- cbind(lmat[, 1] + 1, c(rep(7, (nrow(lmat) -1)/2), 1,rep(NA, (nrow(lmat) -1)/2)), lmat[, 2] + 1)
			}
			else{
				lmat <- cbind(lmat[, 1] + 1, c(rep(NA, nrow(lmat) -1), 1), lmat[, 2] + 1)
			}
									
			lwid <- c(lwid[1], 0.2*length(colnames(RowSideColors)), lwid[2])
		}
		lmat[is.na(lmat)] <- 0
	}
	if (length(lhei) != nrow(lmat)) 
		stop("lhei must have length = nrow(lmat) = ", nrow(lmat))
	if (length(lwid) != ncol(lmat)) 
		stop("lwid must have length = ncol(lmat) =", ncol(lmat))
	op <- par(no.readonly = TRUE)
	on.exit(par(op))
	layout(lmat, widths = lwid, heights = lhei, respect = FALSE)
	if (!is.null(RowSideColors)) {
		par(mar = c(margins[1], 0, 0, 0.5))
		rsc = RowSideColors[rowInd, ]
		rsc.colors = matrix()
		rsc.names = names(table(rsc))
		rsc.i = 1
		for (rsc.name in rsc.names) {
			rsc.colors[rsc.i] = rsc.name
			rsc[rsc == rsc.name] = rsc.i
			rsc.i = rsc.i + 1
		}
		rsc = matrix(as.numeric(rsc), nrow = dim(rsc)[1])
		image(t(rsc), col = as.vector(rsc.colors), axes = FALSE)
		if (length(colnames(RowSideColors)) > 0) {
			axis(3, 0:(dim(rsc)[2] - 1)/(dim(rsc)[2] - 1), RowSideColorsName, 
					las = 2, tick = FALSE, font=2)
		}
	}
	if (!is.null(ColSideColors)) {
		par(mar = c(0.5, 0, 0, margins[2]))
		csc = ColSideColors[colInd, ]
		csc.colors = matrix()
		csc.names = names(table(csc))
		csc.i = 1
		for (csc.name in csc.names) {
			csc.colors[csc.i] = csc.name
			csc[csc == csc.name] = csc.i
			csc.i = csc.i + 1
		}
		csc = matrix(as.numeric(csc), nrow = dim(csc)[1])
		image(csc, col = as.vector(csc.colors), axes = FALSE)
		if (length(colnames(ColSideColors)) > 0) {
			axis(2, 0:(dim(csc)[2] - 1)/(dim(csc)[2] - 1), colnames(ColSideColors), 
					las = 2, tick = FALSE)
		}
	}
	par(mar = c(margins[1], 2.0, 0, margins[2]))
	x <- t(x)
	cellnote <- t(cellnote)
	if (revC) {
		iy <- nr:1
		if (exists("ddr")) 
			ddr <- rev(ddr)
		x <- x[, iy]
		cellnote <- cellnote[, iy]
	}
	else iy <- 1:nr
	image(1:nc, 1:nr, x, xlim = 0.5 + c(0, nc), ylim = 0.5 + 
					c(0, nr), axes = FALSE, xlab = "", ylab = "", col = col, 
			breaks = breaks, ...)  ##should use zlim+capUplimit+downLimit or breaks+capUplimit+downLimit to control the up and down-limit of plotting. 
	retval$carpet <- x
	if (exists("ddr")) 
		retval$rowDendrogram <- ddr
	if (exists("ddc")) 
		retval$colDendrogram <- ddc
	retval$breaks <- breaks
	retval$col <- col
	if (!invalid(na.color) & any(is.na(x))) {
		mmat <- ifelse(is.na(x), 1, NA)
		image(1:nc, 1:nr, mmat, axes = FALSE, xlab = "", ylab = "", ##this make missing value too significant sometimes, since it add na.col later..
				col = na.color, add = TRUE)
	}
	axis(1, 1:nc, labels = labCol, las = 2, line = -0.5, tick = 0, 
			cex.axis = cexCol)
	if (!is.null(xlab)) 
		mtext(xlab, side = 1, line = margins[1] - 1.25)
	axis(4, iy, labels = labRow, las = 2, line = -0.5, tick = 0, 
			cex.axis = cexRow)
	if (!is.null(ylab)) 
		mtext(ylab, side = 4, line = margins[2] - 1.25)
	if (!missing(add.expr)) 
		eval(substitute(add.expr))
	if (!missing(colsep)) 
		for (csep in colsep) rect(xleft = csep + 0.5, ybottom = rep(0, 
							length(csep)), xright = csep + 0.5 + sepwidth[1], 
					ytop = rep(ncol(x) + 1, csep), lty = 1, lwd = 1, 
					col = sepcolor, border = sepcolor)
	if (!missing(rowsep)) 
		for (rsep in rowsep) rect(xleft = 0, ybottom = (ncol(x) + 
								1 - rsep) - 0.5, xright = nrow(x) + 1, ytop = (ncol(x) + 
								1 - rsep) - 0.5 - sepwidth[2], lty = 1, lwd = 1, 
					col = sepcolor, border = sepcolor)
	min.scale <- min(breaks)
	max.scale <- max(breaks)
	x.scaled <- scale01(t(x), min.scale, max.scale)
	if (trace %in% c("both", "column")) {
		retval$vline <- vline
		vline.vals <- scale01(vline, min.scale, max.scale)
		for (i in colInd) {
			if (!is.null(vline)) {
				#abline(v = i - 0.5 + vline.vals, col = linecol, 
						#lty = 2)
			abline(v = vline, col = linecol, 
					lty = 2)
			}
			xv <- rep(i, nrow(x.scaled)) + x.scaled[, i] - 0.5
			xv <- c(xv[1], xv)
			yv <- 1:length(xv) - 0.5
			lines(x = xv, y = yv, lwd = 1, col = tracecol, type = "s")
		}
	}
	if (trace %in% c("both", "row")) {
		retval$hline <- hline
		hline.vals <- scale01(hline, min.scale, max.scale)
		for (i in rowInd) {
			if (!is.null(hline)) {
				abline(h = i + hline, col = linecol, lty = 2)
			}
			yv <- rep(i, ncol(x.scaled)) + x.scaled[i, ] - 0.5
			yv <- rev(c(yv[1], yv))
			xv <- length(yv):1 - 0.5
			lines(x = xv, y = yv, lwd = 1, col = tracecol, type = "s")
		}
	}
	if (!missing(cellnote)) 
		text(x = c(row(cellnote)), y = c(col(cellnote)), labels = c(cellnote), 
				col = notecol, cex = notecex)
	par(mar = c(margins[1], 0, 0, 0))
	if (dendrogram %in% c("both", "row")) {
		plot(ddr, horiz = TRUE, axes = FALSE, yaxs = "i", leaflab = "none")
	}
	else plot.new()
	par(mar = c(0, 0, if (!is.null(main)) 5 else 0, margins[2]))
	if (dendrogram %in% c("both", "column")) {
		plot(ddc, axes = FALSE, xaxs = "i", leaflab = "none")
	}
	else plot.new()
	if (!is.null(main)) 
		title(main, cex.main = 1.5 * op[["cex.main"]])
	if (key) {
		par(mar = c(2*keysize, 0.5*keysize, 3*keysize, 0.5*keysize), cex = 0.6)
		tmpbreaks <- breaks
		if (symkey) {
			max.raw <- max(abs(c(x, breaks)), na.rm = TRUE)
			min.raw <- -max.raw
			tmpbreaks[1] <- -max(abs(x), na.rm = TRUE)
			tmpbreaks[length(tmpbreaks)] <- max(abs(x), na.rm = TRUE)
		}
		else {
			min.raw <- min(x, na.rm = TRUE)
			max.raw <- max(x, na.rm = TRUE)
		}
		z <- seq(min.raw, max.raw, length = length(col))
		image(z = matrix(z, ncol = 1), col = col, breaks = tmpbreaks, 
				xaxt = "n", yaxt = "n")##should use zlim+capUplimit+downLimit or breaks+capUplimit+downLimit to control the up and down-limit of plotting.
		par(usr = c(0, 1, 0, 1))
		lv <- pretty(breaks)
		xv <- scale01(as.numeric(lv), min.raw, max.raw)
		axis(1, at = xv, labels = lv)
		if (scale == "row") 
			mtext(side = 1, "Row Z-Score", line = 2)
		else if (scale == "column") 
			mtext(side = 1, "Column Z-Score", line = 2)
		#else mtext(side = 1, "Value", line = 2)
		if (density.info == "density") {
			dens <- density(x, adjust = densadj, na.rm = TRUE)
			omit <- dens$x < min(breaks) | dens$x > max(breaks)
			dens$x <- dens$x[-omit]
			dens$y <- dens$y[-omit]
			dens$x <- scale01(dens$x, min.raw, max.raw)
			lines(dens$x, dens$y/max(dens$y) * 0.95, col = denscol, 
					lwd = 1)
			axis(2, at = pretty(dens$y)/max(dens$y) * 0.95, pretty(dens$y))
			#title("Color Key\nand Density Plot")
			par(cex = 0.5)
			mtext(side = 2, "Density", line = 2)
		}
		else if (density.info == "histogram") {
			h <- hist(x, plot = FALSE, breaks = breaks)
			hx <- scale01(breaks, min.raw, max.raw)
			hy <- c(h$counts, h$counts[length(h$counts)])
			lines(hx, hy/max(hy) * 0.95, lwd = 1, type = "s", 
					col = denscol)
			axis(2, at = pretty(hy)/max(hy) * 0.95, pretty(hy))
			#title("Color Key\nand Histogram")
			par(cex = 0.5)
			mtext(side = 2, "Count", line = 2)
		}
		else title(keyName,cex=0.5, font=1)
	}
	else plot.new()
	#addAverage = NULL, averageData=NULL,clusterNums=NULL, subClusterOrder =NULL, content=NULL, colAveragePlot=NULL,xAxisSeqForAvePlot=NULL, yAxisSeqForAvePlot=NULL,axisSeq=NULL,

	if(addAverage){
		

		for(clusterNum in c(1:clusterNums)){
			cluster_name<-names(clusterOrder[clusterOrder==clusterNum])
			valueTmpCluster<-averageData[cluster_name,]
			valueGch1<-array()

			for(j in dataSeq){	
				#valueGch1<-cbind(valueGch1,mean(colMeans(valueTmpCluster[,(j-dataStep/2):(j+dataStep/2-1)], na.rm=T), na.rm=T))
				if(as.numeric(floor(j+dataStep/(2*bin_size_in_align)-1)-ceiling(j-dataStep/(2*bin_size_in_align))) <= 0){
					valueGch1<-cbind(valueGch1,mean(valueTmpCluster[,ceiling(j-dataStep/(2*bin_size_in_align))], na.rm=T))	
				}else{
					#tmp<-valueTmpCluster[,ceiling(i-dataStep/(2*bin_size_align)):floor(i+dataStep/(2*bin_size_align)-1)]
					#x<-ifelse(colSums(is.na(tmp))<length(tmp[,1])-1,colMeans(tmp,na.rm=T),NA)
					#valueGch1<-cbind(valueGch1,mean(x,na.rm=T))
					valueGch1<-cbind(valueGch1,mean(colMeans(valueTmpCluster[,ceiling(j-dataStep/(2*bin_size_in_align)):floor(j+dataStep/(2*bin_size_in_align)-1)], na.rm=T), na.rm=T))	
				}
			}
			valueGch1<-valueGch1[,2:length(valueGch1[1,])]	
			par(mar = c(2.0,1.1,0,0))
			if(logScale){
				valueGch1<-log2(valueGch1)
			}
			plot(axisSeq,valueGch1,type="l",axes=F,ylim=c(y_min,y_max), xlab="",ylab="",xaxt='n',ann=F,yaxt='n',col=colAveragePlot[clusterNum],lty=1,font=2,lwd=2)
			par(mar = c(2.0,1.1,0,0),new=T)
		}
		
		axis(1,at=xAxisSeqForAvePlot,lty=1,font=2,cex.axis=1.0,cex.lab=0.8,font.lab=1,lwd=2)
		axis(2,at=seq(y_min,y_max,by=(y_max-y_min)/5),lty=1,font=2,cex.axis=1.0,cex.lab=1.0,font.lab=1,lwd=2)
		title(xlab=xlabForAvePlot, ylab=ylabForAvePlot)
		if(averageLegend){
			legend("topright",paste("cluster",1:clusterNums,sep=""),col=colAveragePlot[1:clusterNum],lty=1,cex=1.2,lwd=2)
		}
		
		#title(xlab="Distance to elements (bp)", ylab="Methylation")
		abline(v=0)
		
		
	}
	else plot.new()
	#mtext(side = 2, ylabForAvePlot, line = 2, cex=0.9, adj=ifelse(is.null(RowSideColors), 0.5, 0.2))
	mtext(side = 2, ylabForAvePlot, line = 2, cex=0.9, adj=0.5)
	##label the clusters
#	if (!missing(RowSideColors)) { 
#		par(mar = c(5, 4, 2, 1), cex = 0.75, mrow=c(length(colnames(RowSideColors)),1))
#		
#		for(ii in c(1:length(colnames(RowSideColors)))){
#			for(colLengend in unique(unlist(RowSideColors[,ii]))){
#				
#				rect(par("usr")[1]*ii/10, par("usr")[3]*ii/10, par("usr")[2]*ii/10, par("usr")[4]*ii/10, col ="colLengend")
#				mtext(side = 2, "Count", line = 2)
				#rect(par("usr")[1]*ii/10, par("usr")[3]/10, par("usr")[2]/10, par("usr")[4]/10, col ="colLengend")
#			}
			
#		}
			
			
#	}
	
	#}
	#else plot.new()

	
	#if (!missing(RowSideColors)) {
	#	par(mar = c(5, 4, 2, 1), cex = 0.75)
		
	#	for(i in c(1:length(colnames(RowSideColors)))){
	#		for(j in unique(factor(RowSideColors[,i]))){
	#			image(z = matrix(1, ncol = 1), col = j, xaxt = "n", yaxt = "n")
	#			par(usr = c(0, 1, 0, 1))
	#			mtext(side = 3, colnames(RowSideColors)[i], line = 2)
	#		}
			
	#	}
	#}
	#else plot.new()
	#if (!missing(ColSideColors)) {
	#	par(mar = c(5, 4, 2, 1), cex = 0.75)
	#	for(i in c(1:length(colnames(ColSideColors)))){
	#		for(j in unique(factor(ColSideColors[,i]))){
	#			image(z = matrix(1, ncol = 1), col = j, xaxt = "n", yaxt = "n")
	#			mtext(side = 3, colnames(ColSideColors)[i], line = 2)
	#		}
	#		
	#	}
	#}
	#else plot.new()
	retval$colorTable <- data.frame(low = retval$breaks[-length(retval$breaks)], 
			high = retval$breaks[-1], color = retval$col)
	invisible(retval)
}

